import { Box, Flex, IconButton, useToast } from '@chakra-ui/react';

import { useCallback, useContext, useState } from 'react';
import { AiOutlineClear } from 'react-icons/ai';
import { MdCloudUpload } from 'react-icons/md';

import { ModelType, TrainingContext } from '..';

import ImagesUploader from './ImagesUploader';
import ParametersForm from './ParametersForm';
import UploadImageView from './UploadImageView';

import trainService from '@/api/train.service';
import IAIButton from '@/common/components/IAIButton';
import IAIScrollArea from '@/common/components/IAIScrollArea';
export type TUploadState = 'not_uploaded' | 'uploading' | 'uploaded';
export type FilePreview = (File | Blob) & { preview: string; uniqId: string };

export default function TrainStudio({ type }: { type: ModelType }) {
  const { setModelChange, modelChange, handleOnAdd } =
    useContext(TrainingContext);
  const [files, setFiles] = useState<FilePreview[]>([]);
  const [uploadedImageUniqIds, setUploadedImageUniqIds] = useState<string[]>(
    []
  );
  const [uploadErrorFiles, setUploadErrorFiles] = useState<FilePreview[]>([]);
  const [uploadState, setUploadState] = useState<TUploadState>('not_uploaded');
  const [errorListUploadState, setErrorListUploadState] =
    useState<TUploadState>('not_uploaded');
  const [uploadedImgIds, setUploadedImgIds] = useState<
    { [k: string]: string }[]
  >([]);
  const toast = useToast();

  const controlUploadTrainImages = (
    list: { formData: FormData; originFile: FilePreview }[],
    num: number,
    cb?: () => void
  ) => {
    let limitNum = num;
    let leftToUploadLength = list.length;
    const taskList = list;
    function run() {
      if (!taskList.length) {
        return;
      }
      const max = Math.min(taskList.length, limitNum);
      for (let i = 0; i < max; i++) {
        const { formData, originFile } = taskList.shift() || {};
        if (!formData || !originFile) {
          continue;
        }
        limitNum--;
        trainService
          .uploadTrainImage(formData)
          .then(({ rt, code }) => {
            if (code === 200) {
              setUploadedImageUniqIds((current) => [
                ...current,
                originFile.uniqId,
              ]);
              setUploadedImgIds((current) => [
                ...current,
                { [originFile.uniqId]: rt.id },
              ]);
              // URL.revokeObjectURL(originFile.preview);
            } else {
              setUploadErrorFiles((f) => [...f, originFile]);
            }
          })
          .catch((e) => {
            setUploadErrorFiles((f) => [...f, originFile]);
            console.log(e);
          })
          .finally(() => {
            limitNum++;
            leftToUploadLength--;
            run();

            if (leftToUploadLength === 0 && cb) {
              cb();
            }
          });
      }
    }
    run();
  };

  const handleUpload = useCallback(
    (isErrorList?: boolean) => {
      if (!isErrorList && files.length < 5) {
        toast({
          title: '为了更好的训练效果，请至少上传5张图片',
          duration: 3000,
          isClosable: true,
          position: 'top',
          status: 'error',
        });
        return;
      }
      if (isErrorList) {
        setErrorListUploadState('uploading');
      } else {
        setUploadState('uploading');
      }
      const filesList = isErrorList ? uploadErrorFiles : files;
      const uploadFetchList = Array.from(filesList).map((f) => {
        const formData = new FormData();
        formData.append('file', f, f.name);
        return { formData, originFile: f };
      });
      setUploadErrorFiles(() => []);
      controlUploadTrainImages(uploadFetchList, 6, () => {
        if (isErrorList) {
          setErrorListUploadState('uploaded');
        } else {
          setUploadState('uploaded');
        }
      });
    },
    [files, uploadErrorFiles, toast]
  );

  const handleDeleteImage = (file: FilePreview) => {
    URL.revokeObjectURL(file.preview);
    setFiles(files.filter((f) => f.uniqId !== file.uniqId));

    const deleteImageIdIndex = Object.keys(uploadedImgIds).findIndex(
      (i) => i === file.uniqId
    );
    const templateIds = uploadedImgIds;
    templateIds.splice(deleteImageIdIndex, 1);
    setUploadedImgIds(templateIds);
  };

  const handleClearImages = () => {
    setFiles([]);
    setUploadState('not_uploaded');
    setUploadedImgIds([]);
    setUploadedImageUniqIds([]);
    setUploadErrorFiles([]);
  };

  const handleCreateProject = () => {
    // 清除上传的图片信息
    files.forEach((file) => {
      URL.revokeObjectURL(file.preview);
    });
    handleClearImages();

    handleOnAdd && handleOnAdd();
    setModelChange && setModelChange(modelChange + 1);
  };

  return (
    <Box p="2" mt="4" rounded="md" shadow="md">
      {uploadState === 'not_uploaded' && (
        <ImagesUploader
          uploadedFiles={files}
          handleFileAccept={(newFiles: FilePreview[]) =>
            setFiles((f) => [...f, ...newFiles])
          }
        />
      )}

      {!!files.length && (
        <Flex pt={3} flexDir="column">
          <Flex justify="flex-end">
            <IconButton
              aria-label="clear"
              icon={<AiOutlineClear />}
              size="xs"
              onClick={handleClearImages}
              variant="imageHoverIconButton"
              fontSize={20}
            />
          </Flex>
          <IAIScrollArea h={files.length > 7 ? 272 : 150} scrollbarSize={6}>
            <Flex flexWrap="wrap" justify="center">
              {files.map((file, index) => (
                <UploadImageView
                  key={file.name + index}
                  file={file}
                  index={index}
                  errorListUploadState={errorListUploadState}
                  uploadState={uploadState}
                  isError={
                    !!uploadErrorFiles.find((f) => f.uniqId === file.uniqId)
                  }
                  isUploaded={uploadedImageUniqIds.indexOf(file.uniqId) >= 0}
                  deleteimage={() => handleDeleteImage(file)}
                />
              ))}
            </Flex>
          </IAIScrollArea>
        </Flex>
      )}

      {files.length > 0 && uploadState !== 'uploaded' && (
        <Flex my={10} textAlign="center" justify="center">
          <Box maxW="20rem">
            <IAIButton
              rightIcon={<MdCloudUpload />}
              size="lg"
              isLoading={uploadState === 'uploading'}
              isDisabled={files.length === 0}
              onClick={() => handleUpload()}
            >
              {`上传 ${files.length} 张图片`}
            </IAIButton>
          </Box>
        </Flex>
      )}
      {uploadErrorFiles.length > 0 && (
        <Flex my={10} textAlign="center" justify="center">
          <Box maxW="20rem">
            <IAIButton
              rightIcon={<MdCloudUpload />}
              size="lg"
              isLoading={errorListUploadState === 'uploading'}
              isDisabled={uploadErrorFiles.length === 0}
              onClick={() => handleUpload(true)}
            >
              {`重新上传 ${uploadErrorFiles.length} 张图片`}
            </IAIButton>
          </Box>
        </Flex>
      )}

      <ParametersForm
        ModalType={type === 'lora' ? '2' : '1'}
        uploadState={uploadState}
        afterCreateStudio={handleCreateProject}
        uploadedImgIds={uploadedImgIds}
      />
    </Box>
  );
}
